#include"KalmanCalculate.hpp"

// 初始化卡尔曼滤波器的状态向量
// 后验状态向量的起点进行初始化
// (this->stateSize = 18;  //状态向量：[cx,cy,vx,vy,ax,ay, tx,ty,tz,vtx,vty,vtz, rx,ry,rz,vrx,vry,vrz]（
// 初始化目标初始状态
//armors[i].firePoint,p_rvec,p_tvec
void KalmanCalculate::init(const Mat &rvec,const Mat &tvec) {
    // 初始状态：[tx, ty, tz, vx, vy, vz, rx, ry, rz, vrx, vry, vrz]
    if(tvec.empty()){
        this->kf.statePost.at<float>(0) = 0;
        this->kf.statePost.at<float>(1) = 0;
        this->kf.statePost.at<float>(2) = 0;
    }else{
        this->kf.statePost.at<float>(0) = static_cast<float>(tvec.at<double>(0));
        this->kf.statePost.at<float>(1) = static_cast<float>(tvec.at<double>(1));
        this->kf.statePost.at<float>(2) = static_cast<float>(tvec.at<double>(2));
    }

    if(rvec.empty()){
        this->kf.statePost.at<float>(6) = 0;  // rx
        this->kf.statePost.at<float>(7) = 0;  // ry
        this->kf.statePost.at<float>(8) = 0;  // rz
    }else{
        this->kf.statePost.at<float>(6) = static_cast<float>(rvec.at<double>(0));
        this->kf.statePost.at<float>(7) = static_cast<float>(rvec.at<double>(1));
        this->kf.statePost.at<float>(8) = static_cast<float>(rvec.at<double>(2));
    }

    // 速度和角速度初始化为0
    this->kf.statePost.at<float>(3) = 0;  // vx
    this->kf.statePost.at<float>(4) = 0;  // vy
    this->kf.statePost.at<float>(5) = 0;  // vz
    this->kf.statePost.at<float>(9) = 0;  // vrx
    this->kf.statePost.at<float>(10) = 0; // vry
    this->kf.statePost.at<float>(11) = 0; // vrz
}

// // 先验状态
// // 预测下一时刻状态
// Point2f KalmanCalculate::priority() {
//     // kf.predict() 会同时计算先验状态（存在 prediction 中）和先验协方差（存在 kf.errorCovPre 中）；
//     cv::Mat prediction = this->kf.predict();
//     return cv::Point2f(prediction.at<float>(0), prediction.at<float>(1));
// }

// // 获取预测的tvec
// Mat KalmanCalculate::predictTvec() {
//     // 确保已经调用了predict
//     cv::Mat tvec(3, 1, CV_64F);
//     tvec.at<double>(0) = this->kf.statePre.at<float>(6);  // tx
//     tvec.at<double>(1) = this->kf.statePre.at<float>(7);  // ty
//     tvec.at<double>(2) = this->kf.statePre.at<float>(8);  // tz
//     return tvec;
// }

// // 获取预测的rvec
// Mat KalmanCalculate::predictRvec(){
//     // 确保已经调用了predict
//     cv::Mat rvec(3, 1, CV_64F);
//     rvec.at<double>(0) = this->kf.statePre.at<float>(12);  // rx
//     rvec.at<double>(1) = this->kf.statePre.at<float>(13);  // ry
//     rvec.at<double>(2) = this->kf.statePre.at<float>(14);  // rz
//     return rvec;
// }

// 获取滤波后的tvec
Mat KalmanCalculate::getTvec(){
    cv::Mat tvec(3, 1, CV_64F);
    tvec.at<double>(0) = static_cast<double>(this->kf.statePost.at<float>(0));  // tx
    tvec.at<double>(1) = static_cast<double>(this->kf.statePost.at<float>(1));  // ty
    tvec.at<double>(2) = static_cast<double>(this->kf.statePost.at<float>(2));  // tz
    return tvec;
}

// 获取滤波后的rvec
Mat KalmanCalculate::getRvec(){
    cv::Mat rvec(3, 1, CV_64F);
    rvec.at<double>(0) = static_cast<double>(this->kf.statePost.at<float>(6));  // rx
    rvec.at<double>(1) = static_cast<double>(this->kf.statePost.at<float>(7));  // ry
    rvec.at<double>(2) = static_cast<double>(this->kf.statePost.at<float>(8));  // rz
    return rvec;
}



// 用当前观测值修正先验状态，最终得到修正后的最优状态（后验状态）
// 用观测值更新滤波状态
void KalmanCalculate::correct(const cv::Mat& rvec,const cv::Mat& tvec) {
    cv::Mat measurement(measSize, 1, CV_32F);
    // 设置tvec观测
    if (tvec.empty()) {
        // 如果没有tvec观测，使用预测值
        measurement.at<float>(0) = this->kf.statePre.at<float>(0);
        measurement.at<float>(1) = this->kf.statePre.at<float>(1);
        measurement.at<float>(2) = this->kf.statePre.at<float>(2);
    } else {
        measurement.at<float>(0) = tvec.at<double>(0);  // 观测到的tx
        measurement.at<float>(1) = tvec.at<double>(1);  // 观测到的ty
        measurement.at<float>(2) = tvec.at<double>(2);  // 观测到的tz
    }
    
    // 设置rvec观测
    if (rvec.empty()) {
        // 如果没有rvec观测，使用预测值
        measurement.at<float>(3) = this->kf.statePre.at<float>(6);
        measurement.at<float>(4) = this->kf.statePre.at<float>(7);
        measurement.at<float>(5) = this->kf.statePre.at<float>(8);
    } else {
        measurement.at<float>(3) = rvec.at<double>(0);  // 观测到的rx
        measurement.at<float>(4) = rvec.at<double>(1);  // 观测到的ry
        measurement.at<float>(5) = rvec.at<double>(2);  // 观测到的rz
    }

    this->kf.correct(measurement);
}

/*

// 设置过程噪声矩阵
void KalmanTracker::setProcessNoise(float dt) {
    float pos_noise = 0.1f * dt;
    float vel_noise = 0.5f * dt;
    float acc_noise = 1.0f * dt;

    cv::setIdentity(kf.processNoiseCov, cv::Scalar::all(1));
    kf.processNoiseCov.at<float>(0,0) = pos_noise;
    kf.processNoiseCov.at<float>(1,1) = pos_noise;
    kf.processNoiseCov.at<float>(2,2) = vel_noise; 
    kf.processNoiseCov.at<float>(3,3) = vel_noise;
    kf.processNoiseCov.at<float>(4,4) = acc_noise;
    kf.processNoiseCov.at<float>(5,5) = acc_noise;
}

*/

// 更新卡尔慢
void KalmanCalculate::updateKalmanTracker(float time_step){
    // 更新参数每帧之间的时间
    this->dt = time_step;

    // 更新状态转移矩阵
    setIdentity(this->kf.transitionMatrix);
    // 平移向量动态模型
    this->kf.transitionMatrix.at<float>(0, 3) = dt;
    this->kf.transitionMatrix.at<float>(1, 4) = dt;  
    this->kf.transitionMatrix.at<float>(2, 5) = dt;
    // 旋转向量动态模型
    this->kf.transitionMatrix.at<float>(6, 9) = dt;
    this->kf.transitionMatrix.at<float>(7, 10) = dt;
    this->kf.transitionMatrix.at<float>(8, 11) = dt;
}

// 获取tvec速度
Mat KalmanCalculate::getTVelocity() {
    cv::Mat tVel(3, 1, CV_64F);
    tVel.at<double>(0) = this->kf.statePost.at<float>(3);  // vx
    tVel.at<double>(1) = this->kf.statePost.at<float>(4);  // vy
    tVel.at<double>(2) = this->kf.statePost.at<float>(5);  // vz
    return tVel;
}

// 获取rvec速度 
Mat KalmanCalculate::getRVelocity() {
    cv::Mat rVel(3, 1, CV_64F);
    rVel.at<double>(0) = this->kf.statePost.at<float>(9);  // vrx
    rVel.at<double>(1) = this->kf.statePost.at<float>(10); // vry
    rVel.at<double>(2) = this->kf.statePost.at<float>(11); // vrz
    return rVel;
}